package com.swsc.ai.tasks;

import com.swsc.ai.util.DateUtil;
import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;

import java.time.LocalDate;
import java.time.format.DateTimeFormatter;

/**
 * @describe: 基金行情相似性计算
 * @author: DuanCXin
 * @created: 2023-10-20 15:52
 */
public class FundPerformSimTask extends TemplateTask {

    @Override
    public String genInputPath(String dataSource) {
        return null;
    }

    @Override
    protected String genSelectColumns() {
        return "innercode, tradingday, unitnv";
    }

    @Override
    public String genWhereCondition(String taskDate) {
        LocalDate localDate = LocalDate.parse(taskDate, DateTimeFormatter.ofPattern("yyyyMMdd"));
        LocalDate work = DateUtil.getWork(localDate.minusDays(1));
        String whereCondition = "where updatestate <= 1 and tradingday <= '"
                + taskDate
                + "' and tradingday >= '"
                + work.format(DateTimeFormatter.ofPattern("yyyyMMdd"))
                + "'";
        return whereCondition;
    }

    @Override
    public void processData(SparkSession session, String outputPath, String... inputParams) {
        // 使用Spark SQL执行查询
        Dataset<Row> sqlDF = session.sql("select tagCode as tarAsset, any_value(rePrices) as rePrices, concat_ws('\t',collect_set(simCode)) AS simAssets, any_value(tradingday) as tradingday from (    \n" +
                "    select tagCode, rePrices, simCode, sub, tradingday, ROW_NUMBER() OVER(PARTITION BY tagCode, tradingday ORDER BY sub ASC) AS rn from (\n" +
                "        select ti.innercode as tagCode, ti.rePrices, tj.innercode as simCode, ABS(ti.rePrices - tj.rePrices) as sub, ti.tradingday from (\n" +
                "            select nv_1.innercode,nv_1.tradingday, nv.tradingday as td, (nv_1.unitnv - nv.unitnv) / nv.unitnv as rePrices from (\n" +
                "                select\n" +
                "                    innercode,\n" +
                "                    tradingday,\n" +
                "                    unitnv,\n" +
                "                    ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingday ASC) AS rn\n" +
                "                from " + inputParams[0] + "\n" +
                "            ) as nv\n" +
                "            inner join (\n" +
                "                select\n" +
                "                    innercode,\n" +
                "                    tradingday,\n" +
                "                    unitnv,\n" +
                "                    ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingday ASC) AS rn\n" +
                "                from " + inputParams[0] + "\n" +
                "            ) as nv_1\n" +
                "            on nv.rn = nv_1.rn - 1 and nv.innercode = nv_1.innercode order by nv.tradingday asc\n" +
                "        ) as ti\n" +
                "        inner join (\n" +
                "            select nv_1.innercode,nv_1.tradingday, nv.tradingday as td, (nv_1.unitnv - nv.unitnv) / nv.unitnv as rePrices from (\n" +
                "              select\n" +
                "                  innercode,\n" +
                "                  tradingday,\n" +
                "                  unitnv,\n" +
                "                  ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingday ASC) AS rn\n" +
                "               from " + inputParams[0] + "\n" +
                "               ) as nv\n" +
                "              inner join (\n" +
                "                select\n" +
                "                    innercode,\n" +
                "                    tradingday,\n" +
                "                    unitnv,\n" +
                "                    ROW_NUMBER() OVER(PARTITION BY innercode ORDER BY tradingday ASC) AS rn\n" +
                "                from " + inputParams[0] + "\n" +
                "            ) as nv_1\n" +
                "         on nv.rn = nv_1.rn - 1 and nv.innercode = nv_1.innercode order by nv.tradingday asc\n" +
                "        ) as tj\n" +
                "        on ti.innercode != tj.innercode and ti.tradingday = tj.tradingday\n" +
                "        order by ti.tradingday, sub ASC\n" +
                "    )\n" +
                ") as k\n" +
                "where k.rn <= 5 \n" +
                "group by tagCode, tradingday\n" +
                "order by tradingday");
       sqlDF.createOrReplaceTempView("tempFundResultView");

        /**
         * 根据正态分布，只取前25%
         */
        Dataset<Row> fundResultPref = session.sql("SELECT s.tarAsset, s.simAssets, s.tradingday FROM (\n" +
                "   SELECT *,  \n" +
                "          RANK() OVER (ORDER BY rePrices) AS rowNum  \n" +
                "   FROM (  \n" +
                "     SELECT *  \n" +
                "     FROM tempFundResultView  \n" +
                "   )\n" +
                ") as s \n" +
                "WHERE s.rowNum <= (SELECT ceil(count(*) / 4) FROM tempFundResultView)");
        /**
         * 根据正态分布，只取后75%
         */
        Dataset<Row> fundResultSuff = session.sql("SELECT s.tarAsset, s.simAssets, s.tradingday FROM (\n" +
                "   SELECT *,  \n" +
                "          RANK() OVER (ORDER BY rePrices) AS rowNum  \n" +
                "   FROM (  \n" +
                "     SELECT *  \n" +
                "     FROM tempFundResultView  \n" +
                "   )\n" +
                ") as s \n" +
                "WHERE s.rowNum >= (SELECT ceil(count(*) * 3 / 4) FROM tempFundResultView)");
        Dataset<Row> resultData = fundResultPref.unionAll(fundResultSuff);
        // 写入HDFS文件
        resultData.write()
                .format("csv")
                .option("header", "true")
                .option("mode", "overwrite")
                .save(outputPath);
    }
}
